#pragma once

#include <memory_allocator_impl.h>
#include <limits>
#include <set>
#include <cassert>
#include <vector>
#include <sstream>

template <std::size_t SmallestSize = 2048, std::size_t BuddyCnt = 16>
class BuddyAllocator : public MemoryPoolImpl
{
public:
    BuddyAllocator() :
        maxChunkSize_(SmallestSize * (1 << BuddyCnt - 1)),
        freeList_(BuddyCnt)
    {
        static_assert(BuddyCnt > 1, "Require BuddyCnt > 1.");
        // 判断SmallestSize为16倍数，且为2的n次幂
        static_assert(SmallestSize >= 16 && (SmallestSize & (SmallestSize - 1)) == 0,
                      "Require SmallestSize >= 16 and is power of 2.");
        // 判断maxChunkSize_溢出
        static_assert(SmallestSize <= (std::numeric_limits<decltype(maxChunkSize_)>::max)() / (1 << (BuddyCnt - 1)),
                      "SmallestSize too large for BuddyCnt");
        char *chunk = (char *)malloc(maxChunkSize_);
        freeList_[BuddyCnt - 1].insert(chunk);
    }

    void *allocate(std::size_t n) override;

    void deallocate(void *const ptr, std::size_t n) override;
    
    std::size_t ROUND_UP(std::size_t n) const override
    {
        return CHUNK_SIZE(findIndex(n));
    }

    ~BuddyAllocator()
    {
    }
private:
    inline std::size_t CHUNK_SIZE(std::size_t index) const
    {
        return SmallestSize * (1 << index);
    }

    inline std::size_t findIndex(std::size_t n) const
    {
        for (std::size_t i = 0; i < freeList_.size(); ++i)
        {
            if (CHUNK_SIZE(i) >= n)
            {
                return i;
            }
        }
        return freeList_.size();
    }

    /// @brief 向上层申请freeList_中索引idx对应大小的块，申请失败则尝试malloc一块新内存
    /// @param idx
    /// @return 指向内存的首地址
    char *chunk_alloc(std::size_t idx);

    /// @brief 将ptr指向的块合并到freeList_中索引idx对应的列表中
    /// @param idx
    /// @param ptr
    /// @return
    char *chunk_merge(std::size_t idx, void *const ptr);

private:
    std::size_t maxChunkSize_;
    std::vector<std::set<char *>> freeList_;
};

template <std::size_t SmallestSize, std::size_t BuddyCnt>
inline void *BuddyAllocator<SmallestSize, BuddyCnt>::allocate(std::size_t n)
{
    if (n <= 0)
        throw std::runtime_error("allocate n require > 0");
    if (n > maxChunkSize_)
    {
        return malloc(n);
    }
    std::size_t idx = findIndex(n);
    assert(0 <= idx && idx < freeList_.size());
    std::set<char *> &chunkSet = freeList_[idx];
    std::size_t bytes = CHUNK_SIZE(idx);
    if (chunkSet.empty()) // chunkSet为空，从上一级buddy分裂出来或malloc一块新的
    {
        char *chunk = chunk_alloc(idx);
        return (void *)chunk;
    }
    char *allocChunk = *chunkSet.begin();
    chunkSet.erase(allocChunk);
    return (void *)(allocChunk);
}

template <std::size_t SmallestSize, std::size_t BuddyCnt>
inline void BuddyAllocator<SmallestSize, BuddyCnt>::deallocate(void *const ptr, std::size_t n)
{
    assert((n & (n - 1)) == 0); // 判断是2的n次幂
    std::size_t idx = findIndex(n);
    if (n > maxChunkSize_)
    {
        free(ptr);
        return;
    }
    chunk_merge(idx, (char *)ptr);
    /*
    std::set<char *> &chunkSet = freeList_[idx];
    std::size_t bytes = CHUNK_SIZE(idx);

    if (bytes < maxChunkSize_) // 当前set中的size小于maxSize，尝试合并
    {
        char *preChunk = (char *)ptr - bytes;
        if (chunkSet.find(preChunk) != chunkSet.end())
        {
            chunkSet.erase(preChunk);
            chunk_merge(idx + 1, preChunk);
            return;
        }

        char *nextChunk = (char *)ptr + bytes;
        if (chunkSet.find(nextChunk) != chunkSet.end())
        {
            chunkSet.erase(nextChunk);
            chunk_merge(idx + 1, (char *)ptr);
            return;
        }
    }
    // 将ptr插入chunkSet
    chunkSet.emplace((char *)ptr);
    */
}

template <std::size_t SmallestSize, std::size_t BuddyCnt>
inline char *BuddyAllocator<SmallestSize, BuddyCnt>::chunk_alloc(std::size_t idx)
{
    assert(0 <= idx && idx < freeList_.size());
    std::size_t allocIdx = idx + 1, chunkSize = 0;
    char *allocChunk = nullptr;
    for (; allocIdx < freeList_.size(); ++allocIdx)
    {
        if (!freeList_[allocIdx].empty())
        {
            auto chunkIt = freeList_[allocIdx].begin();
            allocChunk = *chunkIt;
            freeList_[allocIdx].erase(chunkIt); // erase后迭代器失效
            chunkSize = CHUNK_SIZE(allocIdx);
            break;
        }
    }
    if (!allocChunk)
    {
        allocChunk = (char *)malloc(maxChunkSize_);
        if (!allocChunk)
        {
            std::stringstream ss;
            ss << __FUNCTION__ << " : " << __LINE__ << " malloc " << maxChunkSize_ << " bytes fail!";
            std::string errmsg = ss.str();
            throw std::runtime_error(errmsg);
        }
        chunkSize = maxChunkSize_;
        allocIdx = freeList_.size() - 1;
    }
    if (idx == allocIdx)
    { // 分配的是最大块，不进行块切分
        return allocChunk;
    }
    // 切分块
    for (int i = allocIdx - 1; i >= idx; --i)
    {
        std::set<char *> &chunkSet = freeList_[i];
        char *newChunk = (char *)allocChunk + CHUNK_SIZE(i);
        chunkSet.emplace(newChunk);
        if (i == 0) break; // idx 是std::size_t类型，与-1比较时会将-1变成无符号类型导致bug
    }
    return allocChunk;
}

template <std::size_t SmallestSize, std::size_t BuddyCnt>
inline char *BuddyAllocator<SmallestSize, BuddyCnt>::chunk_merge(std::size_t idx, void *const ptr)
{
    assert(0 <= idx && idx < freeList_.size());
    std::set<char *> &chunkSet = freeList_[idx];
    std::size_t bytes = CHUNK_SIZE(idx);
    if (bytes < maxChunkSize_) // 当前set中的size小于maxSize，尝试合并
    {
        char *preChunk = (char *)ptr - bytes;
        if (chunkSet.find(preChunk) != chunkSet.end())
        {
            chunkSet.erase(preChunk);
            chunk_merge(idx + 1, preChunk);
            return (char *)ptr;
        }

        char *nextChunk = (char *)ptr + bytes;
        if (chunkSet.find(nextChunk) != chunkSet.end())
        {
            chunkSet.erase(nextChunk);
            chunk_merge(idx + 1, (char *)ptr);
            return (char *)ptr;
        }
    }
    chunkSet.emplace((char *)ptr);
    return (char *)ptr;
}
